Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

tests/gdn/test_decode_delta_rule.py (1)

2182-2193: ⚠️ Potential issue | 🟡 Minor

Use exact equality for untouched slots.

These assertions are checking for no mutation, not closeness. Keeping atol/rtol=1e-3 can hide a small stray write in the read slots or the untouched remainder.

Tighten the unchanged-slot checks

     # Read slots must be unchanged (we wrote to different slots)
     torch.testing.assert_close(
-        pool_under_test[read_indices], pool_orig[read_indices], atol=atol, rtol=rtol
+        pool_under_test[read_indices],
+        pool_orig[read_indices],
+        atol=0.0,
+        rtol=0.0,
     )
 
     # Other slots must be unchanged
     used_mask = torch.zeros(pool_size, dtype=torch.bool, device=device)
     used_mask[read_indices] = True
     used_mask[write_indices] = True
     torch.testing.assert_close(
-        pool_under_test[~used_mask], pool_orig[~used_mask], atol=atol, rtol=rtol
+        pool_under_test[~used_mask],
+        pool_orig[~used_mask],
+        atol=0.0,
+        rtol=0.0,
     )

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/gdn/test_decode_delta_rule.py` around lines 2182 - 2193, Change the
"unchanged slots" checks to require exact equality instead of approximate
closeness: for the read-only check comparing pool_under_test[read_indices] and
pool_orig[read_indices] and for the unused-slot check comparing
pool_under_test[~used_mask] and pool_orig[~used_mask], replace the current
torch.testing.assert_close usage (which uses atol/rtol) with an exact-equality
assertion (e.g., use torch.testing.assert_close with rtol=0 and atol=0 or
torch.testing.assert_equal / torch.equal) so any stray writes to pool_under_test
(referenced by pool_under_test, pool_orig, read_indices, write_indices,
used_mask) will fail the test.

flashinfer/gdn_kernels/gdn_decode_bf16_state.py (1)

2549-2553: ⚠️ Potential issue | 🔴 Critical

Preserve null-buffer semantics when defaulting write indices.

When output_state_indices is omitted, this copies -1 padding markers from initial_state_indices onto the write side. The kernel clamps cache_idx, but final writeback still uses write_cache_idx, so padded rows can store before the pool instead of slot 0.

Minimal fix

-    if output_state_indices is None:
-        output_state_indices = initial_state_indices
-    elif output_state_indices.dtype != torch.int32:
-        output_state_indices = output_state_indices.to(torch.int32)
+    if output_state_indices is None:
+        # Preserve the slot-0 null-buffer behavior for padded rows.
+        output_state_indices = initial_state_indices.clamp_min(0)
+    if output_state_indices.dtype != torch.int32:
+        output_state_indices = output_state_indices.to(torch.int32)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@flashinfer/gdn_kernels/gdn_decode_bf16_state.py` around lines 2549 - 2553,
The bug is that defaulting output_state_indices to initial_state_indices copies
-1 padding markers to the write side, allowing padded rows to be written;
instead, when output_state_indices is None set it to a new int32 tensor of the
same shape filled with -1 (preserving null-buffer/write-disabled semantics)
rather than aliasing initial_state_indices; ensure subsequent dtype logic still
converts tensors to torch.int32 (use torch.full_like(initial_state_indices, -1,
dtype=torch.int32) or torch.full(initial_state_indices.shape, -1,
dtype=torch.int32) for the assignment to output_state_indices).

flashinfer/gdn_decode.py (2)

199-210: ⚠️ Potential issue | 🔴 Critical

Reject unsafe pool remaps.

output_state_indices still allows duplicate destinations and cross-batch read/write overlap. In the in-place pool path that makes the final state CTA-order dependent instead of equivalent to gather→compute→scatter. Please reject those mappings here, or route remaps through a staged buffer.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@flashinfer/gdn_decode.py` around lines 199 - 210, The current
output_state_indices validation allows duplicate targets and cross-batch
overlaps which makes the in-place pool path CTA-order dependent; update the
check in the block that currently validates output_state_indices (the code
around the output_state_indices assertions in gdn_decode.py) to reject any
non-permutation mapping: assert that all values are in range [0, B) and that
torch.unique(output_state_indices).numel() == B (i.e., no duplicates and full
bijection), and raise a clear assertion/error message like "output_state_indices
must be a permutation (no duplicates or cross-batch overlaps) when using pool
mode"; alternatively if you prefer to support non-permutations implement a
staged-buffer path (copy to a temporary buffer then scatter) and route
non-permutation remaps through that path instead of the in-place pool path.

---

199-210: ⚠️ Potential issue | 🔴 Critical

Fail fast on invalid destination slots.

This still only checks shape/dtype. A CPU tensor, a negative destination, or an index >= pool_size can reach the kernels and turn into an invalid state write.

Suggested guard

     if output_state_indices is not None:
         assert use_pool, (
             "output_state_indices can only be used with initial_state (pool mode)"
         )
         assert output_state_indices.shape == (B,), (
             f"Expected output_state_indices shape [{B}], "
             f"got {output_state_indices.shape}"
         )
         assert output_state_indices.dtype in (torch.int32, torch.int64), (
             f"output_state_indices must be int32 or int64, "
             f"got {output_state_indices.dtype}"
         )
+        assert output_state_indices.device == initial_state.device, (
+            "output_state_indices must be on the same device as initial_state"
+        )
+        pool_size = int(initial_state.shape[0])
+        in_range = (output_state_indices >= 0) & (
+            output_state_indices < pool_size
+        )
+        assert torch.all(in_range).item(), (
+            f"output_state_indices must be in [0, {pool_size})"
+        )

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@flashinfer/gdn_decode.py` around lines 199 - 210, Add strict runtime guards
for output_state_indices: when output_state_indices is not None (and use_pool is
true), assert it is on the same device as the pool/initial_state (or at least a
CUDA device, not CPU), and assert all indices are within [0, pool_size-1] (no
negatives and none >= pool_size) before passing to kernels; use tensor
operations like output_state_indices.min() and output_state_indices.max() (or
torch.any checks) to detect out-of-range values and raise clear AssertionError
messages referencing output_state_indices, use_pool, and pool_size so invalid
destination slots cannot reach the GPU kernels.

tests/gdn/test_decode_delta_rule.py (1)

2182-2193: ⚠️ Potential issue | 🟡 Minor

Use exact equality for slots that must remain untouched.

These assertions verify that unmodified pool slots remain unchanged. Using atol=1e-3, rtol=1e-3 can mask small stray writes. For immutability checks, use zero tolerance.

💡 Tighten the unchanged-slot assertions

     # Read slots must be unchanged (we wrote to different slots)
     torch.testing.assert_close(
-        pool_under_test[read_indices], pool_orig[read_indices], atol=atol, rtol=rtol
+        pool_under_test[read_indices],
+        pool_orig[read_indices],
+        atol=0.0,
+        rtol=0.0,
     )

     # Other slots must be unchanged
     used_mask = torch.zeros(pool_size, dtype=torch.bool, device=device)
     used_mask[read_indices] = True
     used_mask[write_indices] = True
     torch.testing.assert_close(
-        pool_under_test[~used_mask], pool_orig[~used_mask], atol=atol, rtol=rtol
+        pool_under_test[~used_mask],
+        pool_orig[~used_mask],
+        atol=0.0,
+        rtol=0.0,
     )

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/gdn/test_decode_delta_rule.py` around lines 2182 - 2193, The assertions
that verify unmodified slots use non-zero tolerances which can hide stray
writes; change the checks on pool_under_test vs pool_orig for read_indices and
~used_mask to use exact equality (e.g., torch.equal or
torch.testing.assert_close with atol=0, rtol=0) so that
pool_under_test[read_indices] and pool_under_test[~used_mask] must match
pool_orig exactly; update the two assert calls referencing pool_under_test,
pool_orig, read_indices, write_indices, and used_mask accordingly.

flashinfer/gdn_kernels/gdn_decode_bf16_state.py (1)

2549-2553: ⚠️ Potential issue | 🔴 Critical

Preserve padding/null-buffer semantics when defaulting output_state_indices.

When initial_state_indices contains -1 (padding markers), directly assigning it to output_state_indices will cause the kernel to compute flat_write_idx = -1 * HV + i_hv (a negative offset) and write to invalid memory at lines 1981-2012. The read path has protection (lines 1227-1229 clamp negative cache_idx to 0), but the write path has no equivalent guard.

🐛 Proposed fix: clamp negative indices to slot 0 for writes

     # Default output indices to read indices
     if output_state_indices is None:
-        output_state_indices = initial_state_indices
-    elif output_state_indices.dtype != torch.int32:
+        # Preserve the existing slot-0 null-buffer behavior for padded rows.
+        output_state_indices = initial_state_indices.clamp(min=0)
+    if output_state_indices.dtype != torch.int32:
         output_state_indices = output_state_indices.to(torch.int32)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@flashinfer/gdn_kernels/gdn_decode_bf16_state.py` around lines 2549 - 2553,
When defaulting output_state_indices to initial_state_indices, avoid assigning
the same tensor with -1 padding because the write path computes flat_write_idx
and will write to negative offsets; instead, in the branch where
output_state_indices is None, create a clone of initial_state_indices, replace
negative values (e.g. -1) with 0 to preserve the null-slot semantics for writes,
and then ensure the tensor is converted to torch.int32 (matching the existing
dtype-check branch). Update the code that sets output_state_indices so it uses
output_state_indices = initial_state_indices.clone();
output_state_indices[output_state_indices < 0] = 0; output_state_indices =
output_state_indices.to(torch.int32) (or equivalent) so flat_write_idx cannot be
negative when used with HV and i_hv.

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes